// Copyright 2016 Google LLC. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

#include <packager/media/codecs/video_slice_header_parser.h>

#include <gtest/gtest.h>

namespace shaka {
namespace media {

TEST(H264VideoSliceHeaderParserTest, BasicSupport) {
  // Taken from bear-640x360.mp4 (video)
  const uint8_t kExtraData[] = {
    // Header
    0x01, 0x64, 0x00, 0x1e, 0xff,
    // SPS count (ignore top three bits)
    0xe1,
    // SPS
    0x00, 0x19,  // Size
    0x67, 0x64, 0x00, 0x1e, 0xac, 0xd9, 0x40, 0xa0,
    0x2f, 0xf9, 0x70, 0x11, 0x00, 0x00, 0x03, 0x03,
    0xe9, 0x00, 0x00, 0xea, 0x60, 0x0f, 0x16, 0x2d,
    0x96,
    // PPS count
    0x01,
    // PPS
    0x00, 0x06,  // Size
    0x68, 0xeb, 0xe3, 0xcb, 0x22, 0xc0
  };
  const uint8_t kData[] = {
    // Incomplete data, but we only care about the header size.
    0x65, 0x88, 0x84, 0x00, 0x21, 0xff, 0xcf, 0x73, 0xc7, 0x24,
    0xc8, 0xc3, 0xa5, 0xcb, 0x77, 0x60, 0x50, 0x85, 0xd9, 0xfc
  };
  const std::vector<uint8_t> extra_data(kExtraData,
                                        kExtraData + std::size(kExtraData));

  H264VideoSliceHeaderParser parser;
  ASSERT_TRUE(parser.Initialize(extra_data));

  Nalu nalu;
  ASSERT_TRUE(nalu.Initialize(Nalu::kH264, kData, std::size(kData)));
  // Real header size is 34 bits, but we round up to 5 bytes.
  EXPECT_EQ(5, parser.GetHeaderSize(nalu));
}

TEST(H264VideoSliceHeaderParserTest, SupportsMultipleEntriesInExtraData) {
  const uint8_t kExtraData[] = {
    // Header
    0x01, 0xed, 0xf0, 0x0d, 0x00,
    // SPS count (ignore top three bits)
    0xe3,
    // SPS
    0x00, 0x19,  // Size
    0x67, 0x64, 0x00, 0x1e, 0xac, 0xd9, 0x40, 0xa0,
    0x2f, 0xf9, 0x70, 0x11, 0x00, 0x00, 0x03, 0x03,
    0xe9, 0x00, 0x00, 0xea, 0x60, 0x0f, 0x16, 0x2d,
    0x96,
    // SPS
    0x00, 0x19,  // Size
    0x67, 0x64, 0x00, 0x1e, 0xac, 0xd9, 0x40, 0xa0,
    0x2f, 0xf9, 0x70, 0x11, 0x00, 0x00, 0x03, 0x03,
    0xe9, 0x00, 0x00, 0xea, 0x60, 0x0f, 0x16, 0x2d,
    0x96,
    // SPS
    0x00, 0x19,  // Size
    0x67, 0x64, 0x00, 0x1e, 0xac, 0xd9, 0x40, 0xa0,
    0x2f, 0xf9, 0x70, 0x11, 0x00, 0x00, 0x03, 0x03,
    0xe9, 0x00, 0x00, 0xea, 0x60, 0x0f, 0x16, 0x2d,
    0x96,
    // PPS count
    0x03,
    // PPS
    0x00, 0x06,  // Size
    0x68, 0xeb, 0xe3, 0xcb, 0x22, 0xc0,
    // PPS
    0x00, 0x06,  // Size
    0x68, 0xeb, 0xe3, 0xcb, 0x22, 0xc0,
    // PPS
    0x00, 0x06,  // Size
    0x68, 0xeb, 0xe3, 0xcb, 0x22, 0xc0
  };
  const std::vector<uint8_t> extra_data(kExtraData,
                                        kExtraData + std::size(kExtraData));

  H264VideoSliceHeaderParser parser;
  EXPECT_TRUE(parser.Initialize(extra_data));
}

TEST(H264VideoSliceHeaderParserTest, IgnoresExtraDataAtEnd) {
  const uint8_t kExtraData[] = {
    // Header
    0x01, 0xed, 0xf0, 0x0d, 0x00,
    // SPS count (ignore top three bits)
    0xe1,
    // SPS
    0x00, 0x19,  // Size
    0x67, 0x64, 0x00, 0x1e, 0xac, 0xd9, 0x40, 0xa0,
    0x2f, 0xf9, 0x70, 0x11, 0x00, 0x00, 0x03, 0x03,
    0xe9, 0x00, 0x00, 0xea, 0x60, 0x0f, 0x16, 0x2d,
    0x96,
    // PPS count
    0x00,
    // Extra data
    0x00, 0x19, 0x67, 0x64, 0x00
  };
  const std::vector<uint8_t> extra_data(kExtraData,
                                        kExtraData + std::size(kExtraData));

  H264VideoSliceHeaderParser parser;
  EXPECT_TRUE(parser.Initialize(extra_data));
}

TEST(H264VideoSliceHeaderParserTest, ErrorsForEOSAfterEntry) {
  const uint8_t kExtraData[] = {
    // Header
    0x01, 0xed, 0xf0, 0x0d, 0x00,
    // SPS count (ignore top three bits)
    0xe3,
    // SPS
    0x00, 0x19,  // Size
    0x67, 0x64, 0x00, 0x1e, 0xac, 0xd9, 0x40, 0xa0,
    0x2f, 0xf9, 0x70, 0x11, 0x00, 0x00, 0x03, 0x03,
    0xe9, 0x00, 0x00, 0xea, 0x60, 0x0f, 0x16, 0x2d,
    0x96,
  };
  const std::vector<uint8_t> extra_data(kExtraData,
                                        kExtraData + std::size(kExtraData));

  H264VideoSliceHeaderParser parser;
  EXPECT_FALSE(parser.Initialize(extra_data));
}

TEST(H264VideoSliceHeaderParserTest, ErrorsForEOSWithinEntry) {
  const uint8_t kExtraData[] = {
    // Header
    0x01, 0xed, 0xf0, 0x0d, 0x00,
    // SPS count (ignore top three bits)
    0xe3,
    // SPS
    0x00, 0x19,  // Size
    0x67, 0x64, 0x00, 0x1e, 0xac, 0xd9, 0x40, 0xa0,
  };
  const std::vector<uint8_t> extra_data(kExtraData,
                                        kExtraData + std::size(kExtraData));

  H264VideoSliceHeaderParser parser;
  EXPECT_FALSE(parser.Initialize(extra_data));
}

TEST(H265VideoSliceHeaderParserTest, BasicSupport) {
  const uint8_t kExtraData[] = {
      // Header
      0x01, 0x01, 0x40, 0x00, 0x00, 0x00, 0x90, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x96, 0xF0, 0x00, 0xFC, 0xFD, 0xF8, 0xF8, 0x00, 0x00, 0x0F, 0x03,
      // VPS
      0x20, 0x00, 0x01, 0x00, 0x17,  // Size = 23
      0x40, 0x01, 0x0C, 0x01, 0xFF, 0xFF, 0x01, 0x40, 0x00, 0x00, 0x03, 0x00,
      0x90, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x96, 0xBC, 0x09,
      // SPS
      0x21, 0x00, 0x01, 0x00, 0x2D,  // Size = 45
      0x42, 0x01, 0x01, 0x01, 0x40, 0x00, 0x00, 0x03, 0x00, 0x90, 0x00, 0x00,
      0x03, 0x00, 0x00, 0x03, 0x00, 0x96, 0xA0, 0x03, 0xC0, 0x80, 0x11, 0x07,
      0xCB, 0x96, 0xF4, 0xA4, 0x21, 0x19, 0x3F, 0x8C, 0x04, 0x04, 0x00, 0x00,
      0x03, 0x00, 0x04, 0x00, 0x00, 0x03, 0x01, 0x68, 0x20,
      // PPS
      0x22, 0x00, 0x01, 0x00, 0x08,  // Size = 8
      0x44, 0x01, 0xC0, 0xAD, 0xF0, 0x33, 0x24, 0x00};
  const uint8_t kDataIntra[] = {
      // Incomplete data, but we only care about the header size.
      0x26, 0x01, 0xAF, 0x2E,
      // Below are just arbitrary data after the slice header.
      0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
  const uint8_t kDataInter[] = {
      // Incomplete data, but we only care about the header size.
      0x02, 0x01, 0xD0, 0x08, 0xBF, 0xA5, 0xC0,
      // Below are just arbitrary data after the slice header.
      0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};

  const std::vector<uint8_t> extra_data(kExtraData,
                                        kExtraData + std::size(kExtraData));
  H265VideoSliceHeaderParser parser;
  ASSERT_TRUE(parser.Initialize(extra_data));

  // Can call IntializeLayered() with empty data in the single layer case.
  EXPECT_TRUE(parser.InitializeLayered(std::vector<uint8_t>()));

  Nalu nalu;
  ASSERT_TRUE(nalu.Initialize(Nalu::kH265, kDataIntra, std::size(kDataIntra)));
  EXPECT_EQ(2, parser.GetHeaderSize(nalu));
  ASSERT_TRUE(nalu.Initialize(Nalu::kH265, kDataInter, std::size(kDataInter)));
  EXPECT_EQ(5, parser.GetHeaderSize(nalu));
}

TEST(H265VideoSliceHeaderParserTest, BasicStereoVideoSupport) {
  const uint8_t kExtraDataL0[] = {
      // Header
      0x01, 0x21, 0x60, 0x00, 0x00, 0x00, 0xB0, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x78, 0xF0, 0x00, 0xFC, 0xFD, 0xF8, 0xF8, 0x00, 0x00, 0x03, 0x04,
      // VPS
      0xA0, 0x00, 0x01, 0x00, 0x37,  // Size = 55
      0x40, 0x01, 0x0C, 0x11, 0xFF, 0xFF, 0x21, 0x60, 0x00, 0x00, 0x03, 0x00,
      0xB0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x78, 0x93, 0x05, 0x6F,
      0x78, 0x20, 0x00, 0x28, 0x24, 0x59, 0x72, 0x60, 0x20, 0x00, 0x00, 0x03,
      0x00, 0xF8, 0x80, 0x00, 0x00, 0x03, 0x00, 0x07, 0x88, 0xD0, 0x78, 0x00,
      0x44, 0x0A, 0x01, 0xE5, 0xC9, 0x7D, 0x20,
      // SPS
      0xA1, 0x00, 0x01, 0x00, 0x28,  // Size = 40
      0x42, 0x01, 0x01, 0x21, 0x60, 0x00, 0x00, 0x03, 0x00, 0xB0, 0x00, 0x00,
      0x03, 0x00, 0x00, 0x03, 0x00, 0x78, 0xA0, 0x03, 0xC0, 0x80, 0x11, 0x07,
      0xCB, 0x96, 0x4E, 0xE4, 0xC9, 0xAE, 0xD4, 0x94, 0x04, 0x04, 0x04, 0x0B,
      0xB4, 0x28, 0x4D, 0x01,
      // PPS
      0xA2, 0x00, 0x01, 0x00, 0x0D,  // Size = 13
      0x44, 0x01, 0xC5, 0xE3, 0x0F, 0x09, 0xC1, 0x80, 0xC7, 0xB0, 0x9A, 0x01,
      0x40,
      // SEI
      0xA7, 0x00, 0x01, 0x00, 0x09,  // Size = 9
      0x4E, 0x01, 0xB0, 0x04, 0x04, 0x0D, 0x00, 0x20, 0x80};
  const uint8_t kExtraDataL1[] = {// Header
                                  0x01, 0xF0, 0x00, 0xFC, 0xCB, 0x02,
                                  // SPS
                                  0xA1, 0x00, 0x01, 0x00, 0x0E,  // Size = 14
                                  0x42, 0x09, 0x0E, 0x85, 0xB9, 0x32, 0x6B,
                                  0xBE, 0x80, 0x2E, 0xD0, 0xA1, 0x34, 0x04,
                                  // PPS
                                  0xA2, 0x00, 0x01, 0x00, 0x0E,  // Size = 14
                                  0x44, 0x09, 0x48, 0x5A, 0x43, 0x0F, 0x09,
                                  0xC1, 0x80, 0xC7, 0xB4, 0x9A, 0x01, 0x40};
  const uint8_t kDataL0Intra[] = {
      // Incomplete data, but we only care about the header size.
      0x26, 0x01, 0xA3, 0xC6, 0xAC, 0x30, 0x02, 0x0C, 0xB0, 0x02, 0x49, 0x88,
      0x7F, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xF8,
      // Below are just arbitrary data after the slice header.
      0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
  const uint8_t kDataL1Intra[] = {
      // Incomplete data, but we only care about the header size.
      0x26, 0x09, 0x90, 0x80, 0x3C, 0xC6, 0xAC, 0x30, 0x01, 0x20, 0x6A, 0x01,
      0x37, 0x28, 0x7F, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xF8,
      // Below are just arbitrary data after the slice header.
      0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
  const uint8_t kDataL0Inter[] = {
      // Incomplete data, but we only care about the header size.
      0x02, 0x01, 0xC4, 0x03, 0x7C, 0xC4, 0xAC, 0x30, 0x00, 0x0F, 0xAA, 0x00,
      0x1B, 0x30, 0x7F, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xF8,
      // Below are just arbitrary data after the slice header.
      0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
  const uint8_t kDataL1Inter[] = {
      // Incomplete data, but we only care about the header size.
      0x02, 0x09, 0xA1, 0x00, 0x9A, 0x00, 0xBD, 0x65, 0x89, 0x58, 0x60, 0x00,
      0x30, 0x2C, 0x00, 0x41, 0xC0, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xF0,
      // Below are just arbitrary data after the slice header.
      0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};

  const std::vector<uint8_t> extra_data_layer_0(
      kExtraDataL0, kExtraDataL0 + std::size(kExtraDataL0));
  H265VideoSliceHeaderParser parser;
  ASSERT_TRUE(parser.Initialize(extra_data_layer_0));
  const std::vector<uint8_t> extra_data_layer_1(
      kExtraDataL1, kExtraDataL1 + std::size(kExtraDataL1));
  ASSERT_TRUE(parser.InitializeLayered(extra_data_layer_1));

  Nalu nalu;
  ASSERT_TRUE(
      nalu.Initialize(Nalu::kH265, kDataL0Intra, std::size(kDataL0Intra)));
  EXPECT_EQ(17, parser.GetHeaderSize(nalu));
  ASSERT_TRUE(
      nalu.Initialize(Nalu::kH265, kDataL1Intra, std::size(kDataL1Intra)));
  EXPECT_EQ(19, parser.GetHeaderSize(nalu));
  ASSERT_TRUE(
      nalu.Initialize(Nalu::kH265, kDataL0Inter, std::size(kDataL0Inter)));
  EXPECT_EQ(19, parser.GetHeaderSize(nalu));
  ASSERT_TRUE(
      nalu.Initialize(Nalu::kH265, kDataL1Inter, std::size(kDataL1Inter)));
  EXPECT_EQ(22, parser.GetHeaderSize(nalu));
}

}  // namespace media
}  // namespace shaka
